#/usr/bin/python
#-*- endcoding:utf-8 -*-

#Author TCY
#Email tangchenyi@mail.hfut.edu.cn
#Date 2015/04/06

import os
import math

class FeatureExtractor:
    '''
    目前只提供TF,TF-IDF两种特征提取
    当指定词典(dictPath)的时候停用词(stopwordPath)无效
    '''
    TYPE_TF = "TF"
    TYPE_TF_IDF = "TF-IDF"
    TYPE = (TYPE_TF, TYPE_TF_IDF)
    
    def __init__(self, featureType, srcDataPath, dstFeaturePath, dictPath=None, stopwordPath=None, wordFilter=None):
        if featureType not in self.TYPE:
            raise ValueError,"featureType error"
        
        self.__featureType = featureType
        self.__srcDataPath = srcDataPath
        self.__dstFeaturePath = dstFeaturePath
        self.__dictPath = dictPath
        self.__stopwordPath = stopwordPath
        if wordFilter is not None:
            self.wordFilter = wordFilter
        else:
            self.wordFilter = self.defaultWordFilter
        pass

    def extract(self):
        self.buildDict()
        if self.__featureType == self.TYPE_TF:
            self.run_TF()
        elif self.__featureType == self.TYPE_TF_IDF:
            self.run_TF_IDF()
        if os.path.exists(self.__dstFeaturePath+".tmp"):
            os.remove(self.__dstFeaturePath+".tmp")
        pass

    def run_TF(self):
        self.count()
        self.write_result()
        pass

    def run_TF_IDF(self):
        global_words_count_dict = self.count()
        self.calculate_IDF_and_write_result(global_words_count_dict)
        pass
    
    def count(self):
        global_words_count_dict = dict()
        fp = file(self.__srcDataPath,"r")
        wfp = file(self.__dstFeaturePath+".tmp","w")
        self.numOfLine = 0.0
        for line in fp:
            words = line.split()
            label = words[0]
            wfp.write(label+" ")
            words_count_dict = dict()
            count = 0.0;
            for word in words[1:]:
                if word in self.__dictionary:
                    if word not in words_count_dict:
                        words_count_dict[word] = 1
                    else:
                        words_count_dict[word] += 1
                    count += 1;
            for word in words_count_dict:
                if word in global_words_count_dict:
                    global_words_count_dict[word]+=1
                else:
                    global_words_count_dict[word] = 1
                TF = words_count_dict[word]/count
                wfp.write(word+":"+str(TF)+" ")
            wfp.write("\n")
            self.numOfLine+=1
        wfp.close()
        fp.close()
        return global_words_count_dict
        pass
        
    def write_result(self):
        fp = file(self.__dstFeaturePath+".tmp","r")
        wfp = file(self.__dstFeaturePath,"w")
        for line in fp:
            words = line.split()
            label = words[0]
            wfp.write(label+" ")
            for word in words[1:]:
                tmp = word.split(":")
                wfp.write("%d:%s "%(self.__dictionary[tmp[0]],tmp[1]))
            wfp.write("\n")
        wfp.close()
        fp.close()
        pass

    def calculate_IDF_and_write_result(self, words_count_dict):
        for word in words_count_dict:
            words_count_dict[word] = math.log(self.numOfLine/words_count_dict[word])
        fp = file(self.__dstFeaturePath+".tmp","r")
        wfp = file(self.__dstFeaturePath,"w")
        for line in fp:
            words = line.split()
            label = words[0]
            wfp.write(label+" ")
            for word in words[1:]:
                tmp = word.split(":")
                wfp.write("%d:%f "%(self.__dictionary[tmp[0]],float(tmp[1])*words_count_dict[tmp[0]]))
            wfp.write("\n")
        wfp.close()
        fp.close()
        pass
                
    def loadDictionary(self, filePath, isDictFile):
        fp=file(filePath,"r")
        self.__dictionary = dict()
        wordId=0
        for line in fp:
            words = line.split()
            if isDictFile is False:
                words = words[1:]
            for word in words:
                if self.wordFilter(word) is False and word not in self.__dictionary:
                    self.__dictionary[word]=wordId
                    wordId+=1
        fp.close()
        pass
    
    def loadStopword(self):
        self.__stopword = set()
        if self.__stopwordPath is None:
            return
        fp=file(self.__stopwordPath,"r")
        
        for line in fp:
            words = line.split()
            for word in words:
                self.__stopword.add(word)
        fp.close()
        pass

    def buildDict(self):
        if self.__dictPath is None:
           self.loadDictionary(self.__srcDataPath, False)
        else:
            self.loadDictionary(self.__dictPath, True)
        if self.__stopwordPath is not None:
            self.loadStopword()
            for w in self.__stopword:
                if w in __dictionary:
                    del self.__dictionary[w]
        pass

    def defaultWordFilter(self,word):
        return False
        pass

def tt(word):
    return False
    pass

if __name__ == "__main__":
    extractor = FeatureExtractor("TF-IDF", "test.txt", "dest.txt")
    extractor.extract()

